import os
import xarray as xr
import numpy as np
import rasterio
from rasterio.transform import from_origin

def calculate_daily_mean_precipitation(input_folder, output_folder):
    # 获取文件夹中的所有.nc文件
    nc_files = [f for f in os.listdir(input_folder) if f.endswith('.nc')]
    
    for nc_file in nc_files:
        # 读取每一个NetCDF文件
        nc_path = os.path.join(input_folder, nc_file)
        ds = xr.open_dataset(nc_path)
        
        if 'SMCI' not in ds.variables:
            print(f"文件 {nc_file} 中未找到变量 'SMCI'")
            continue

        # 计算日均降水量（即对时间维度进行平均）
        daily_mean_precip = ds['SMCI'].mean(dim='time', skipna=True)

        # 将数据类型转换为16位整型，并设置Nodata值为-999
        daily_mean_precip = daily_mean_precip.fillna(-999).astype(np.int16)
        
        # 获取空间元数据信息
        transform = from_origin(ds.lon[0], ds.lat[0], abs(ds.lon[1] - ds.lon[0]), abs(ds.lat[1] - ds.lat[0]))
        width, height = daily_mean_precip.shape

        # 设置输出文件名
        output_file = os.path.join(output_folder, f"{os.path.splitext(nc_file)[0]}_mean.tif")

        # 将结果保存为GeoTIFF
        with rasterio.open(
            output_file,
            'w',
            driver='GTiff',
            height=height,
            width=width,
            count=1,
            dtype=rasterio.int16,  # 设置数据类型为16位整型
            crs='EPSG:4326',       # 设置CRS（可以根据实际需求调整）
            transform=transform,
            nodata=-999            # 设置NoData值为-999
        ) as dst:
            dst.write(daily_mean_precip.values, 1)
        
        print(f"已保存日均降水量TIFF文件: {output_file}")

# 使用示例
input_folder = 'E:\\1km分辨率土壤湿度\\10cm'  # 替换为NetCDF文件所在文件夹路径
output_folder = 'E:\\1km分辨率土壤湿度\\10cm'  # 替换为保存TIFF文件的文件夹路径
calculate_daily_mean_precipitation(input_folder, output_folder)
